#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：Flower_Class 
@File    ：train_model.py
@IDE     ：PyCharm 
@Author  ：Rice_cc
@Date    ：2024/10/16 1:22 
'''

import tensorflow as tf
import matplotlib.pyplot as plt

# import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"


# 数据加载，按照8:2的比例加载花卉数据
def data_load(data_dir, img_height, img_width, batch_size):
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir,
        label_mode='categorical',
        validation_split=0.2,
        subset="training",
        seed=123,
        image_size=(img_height, img_width),
        batch_size=batch_size)
    val_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir,
        label_mode='categorical',
        validation_split=0.2,
        subset="validation",
        seed=123,
        image_size=(img_height, img_width),
        batch_size=batch_size)
    class_names = train_ds.class_names

    return train_ds, val_ds, class_names


# 模型加载，指定图像处理的大小
def model_load(IMG_SHAPE=(224, 224, 3)):
    # 微调的过程中不需要进行归一化处理
    base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
                                                   include_top=False,  # 网络最后是否包含全连接层
                                                   weights='imagenet')  # 预训练权重, imagenet 默认值
    base_model.trainable = False  # 冻结卷积模块
    # 描述神经网络结构
    model = tf.keras.models.Sequential([
        tf.keras.layers.experimental.preprocessing.Rescaling(1. / 127.5, offset=-1,
                                                             input_shape=IMG_SHAPE), base_model,  # 全局均值池化
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(5, activation='softmax')
    ])
    # 输出模型各层的参数状况
    model.summary()
    # 模型训练（优化器，损失函数，准确率）
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    return model


def train(epochs, is_transfer=False):
    train_ds, val_ds, class_names = data_load("C:/Users/23987/Downloads/flower_photos/flower_split/train",
                                              224, 224, 4)
    model = model_load()
    history = model.fit(train_ds, validation_data=val_ds, epochs=epochs)  # 保存训练模型
    model.save("models/mobilenet_flower.h5")
    show_loss_acc(history)


def show_loss_acc(history):
    print(str(history))
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']

    loss = history.history['loss']
    val_loss = history.history['val_loss']

    plt.figure(figsize=(8, 8))
    plt.subplot(2, 1, 1)
    plt.plot(acc, label='Training Accuracy')
    plt.plot(val_acc, label='Validation Accuracy')
    plt.legend(loc='lower right')
    plt.ylabel('Accuracy')
    plt.ylim([min(plt.ylim()), 1])
    plt.title("Training and Validation Accuracy")

    plt.subplot(2, 1, 2)
    plt.plot(loss, label='Training Loss')
    plt.plot(val_loss, label='Validation Loss')
    plt.legend(loc='upper right')
    plt.ylabel('Cross Entropy')
    plt.ylim([0, 1.0])
    plt.title('Training and Validation Loss')
    plt.xlabel('epoch')
    plt.show()


if __name__ == "__main__":
    train(epochs=10)
